#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Piper机械臂Xbox手柄控制系统 V1.0
作者: Yinchuan Li
日期: 2025-07-27
功能: 通过XBox手柄实时控制Piper机械臂的6自由度移动和夹爪操作
版本: 1.0

主要功能:
1. 左摇杆控制X-Y轴移动（30Hz持续控制）
2. LT/RT触发器控制Z轴上下移动（30Hz持续控制）
3. LB/RB按钮控制夹爪开合
4. Menu按钮重置到初始位置
5. 6cm/秒移动速度，30Hz控制频率
"""

import pygame
import time
import sys
from piper_sdk import *

class XboxPiperController:
    """
    Xbox手柄控制Piper机械臂类
    支持6自由度实时控制和夹爪操作
    """
    
    def __init__(self, piper):
        """
        初始化Xbox控制器
        
        Args:
            piper: Piper机械臂接口对象
        """
        self.piper = piper
        
        # 初始化pygame和手柄
        self._init_xbox_controller()
        
        # 机械臂控制参数
        self.factor = 1000  # 单位转换因子（毫米转微米）
        self.move_speed = 6.0  # 移动速度（厘米/秒）
        self.control_frequency = 30.0  # 控制频率（Hz）
        self.move_distance_per_frame = self.move_speed / self.control_frequency  # 每帧移动距离（厘米）
        
        # 当前位置和初始位置
        self.current_position = None  # 当前机械臂位置 [X, Y, Z, RX, RY, RZ]
        self.initial_position = None  # 初始位置
        self.safe_initial_pose = [57.0, 0.0, 215.0, 0.0, 85.0, 0.0]  # 安全初始位姿
        
        # 摇杆控制参数
        self.deadzone = 0.1  # 摇杆死区，防止微小抖动
        self.last_joystick_state = {'x': 0, 'y': 0, 'lt': 0, 'rt': 0}  # 记录上一次摇杆状态
        
        # 按钮状态记录
        self.last_button_states = {}
        
        # 夹爪控制参数
        self.gripper_delay = 0.3  # 夹爪操作延时（秒）
        
        # 初始化机械臂
        self._initialize_robot()
        
    def _init_xbox_controller(self):
        """初始化Xbox手柄"""
        pygame.init()
        pygame.joystick.init()
        
        if pygame.joystick.get_count() == 0:
            raise RuntimeError("未检测到Xbox手柄，请确保手柄已连接")
            
        self.joystick = pygame.joystick.Joystick(0)
        self.joystick.init()
        
        print(f"已连接Xbox手柄: {self.joystick.get_name()}")
        print(f"轴数量: {self.joystick.get_numaxes()}, 按钮数量: {self.joystick.get_numbuttons()}")
        
        # Xbox按键映射
        self.button_mapping = {
            4: "LB",      # 左肩键 - 关闭夹爪
            5: "RB",      # 右肩键 - 打开夹爪
            7: "Menu",    # 菜单键 - 重置位置
        }
        
        # 摇杆轴映射
        self.axis_mapping = {
            0: "Left Stick X",    # 左摇杆X轴 - 控制X轴移动
            1: "Left Stick Y",    # 左摇杆Y轴 - 控制Y轴移动
            2: "LT Trigger",      # 左触发器 - Z轴向下
            5: "RT Trigger",      # 右触发器 - Z轴向上
        }
        
    def _initialize_robot(self):
        """初始化机械臂"""
        print("正在初始化机械臂...")
        
        try:
            # 连接机械臂
            print("步骤1: 连接机械臂...")
            self.piper.ConnectPort()
            print("机械臂连接成功")
            
            # 使能机械臂
            print("步骤2: 使能机械臂...")
            self.piper.EnableArm(7)
            print("机械臂使能命令已发送")
            
            # 使能检测
            print("步骤3: 检测使能状态...")
            self._enable_robot()
            
            # 移动到安全初始位姿
            print("步骤4: 移动到安全初始位姿...")
            self._move_to_initial_pose()
            
            # 等待移动完成并获取实际位置
            print("步骤5: 获取当前位置...")
            time.sleep(2)
            self.current_position = self._get_current_position()
            
            # 如果获取位置失败，使用预定义位置
            if self.current_position is None:
                print("警告: 无法获取当前位置，使用预定义位姿")
                self.current_position = self.safe_initial_pose.copy()
            else:
                print("成功获取当前位置")
            
            # 保存初始位置用于重置功能
            self.initial_position = self.current_position.copy()
            print("初始位置已保存")
            
            # 显示初始位置信息
            print(f"初始位置: X={self.current_position[0]:.1f}, Y={self.current_position[1]:.1f}, Z={self.current_position[2]:.1f}")
            print(f"         RX={self.current_position[3]:.1f}, RY={self.current_position[4]:.1f}, RZ={self.current_position[5]:.1f}")
            print(f"控制参数: 移动速度={self.move_speed}cm/s, 控制频率={self.control_frequency}Hz")
            print(f"         每帧移动距离={self.move_distance_per_frame:.3f}cm")
            print("机械臂初始化完成！")
            
        except Exception as e:
            print(f"机械臂初始化失败: {e}")
            import traceback
            traceback.print_exc()
            raise
        
    def _enable_robot(self):
        """使能机械臂并检测使能状态"""
        enable_flag = False
        timeout = 5  # 超时时间（秒）
        start_time = time.time()
        
        while not enable_flag:
            elapsed_time = time.time() - start_time
            print("检查使能状态...")
            
            try:
                # 获取所有电机的使能状态
                low_spd_info = self.piper.GetArmLowSpdInfoMsgs()
                enable_flag = (
                    low_spd_info.motor_1.foc_status.driver_enable_status and
                    low_spd_info.motor_2.foc_status.driver_enable_status and
                    low_spd_info.motor_3.foc_status.driver_enable_status and
                    low_spd_info.motor_4.foc_status.driver_enable_status and
                    low_spd_info.motor_5.foc_status.driver_enable_status and
                    low_spd_info.motor_6.foc_status.driver_enable_status
                )
            except Exception as e:
                print(f"获取使能状态失败: {e}")
                enable_flag = False
            
            # 如果未使能，尝试重新使能
            if not enable_flag:
                self.piper.EnableArm(7)
                self.piper.GripperCtrl(0, 1000, 0x01, 0)
            
            # 超时检查
            if elapsed_time > timeout:
                print("程序自动使能超时,退出程序")
                exit(0)
                
            time.sleep(1)
        print("机械臂使能成功")
    
    def _move_to_initial_pose(self):
        """移动到安全初始位姿"""
        try:
            # 将位置转换为微米单位
            X = round(self.safe_initial_pose[0] * self.factor)
            Y = round(self.safe_initial_pose[1] * self.factor)
            Z = round(self.safe_initial_pose[2] * self.factor)
            RX = round(self.safe_initial_pose[3] * self.factor)
            RY = round(self.safe_initial_pose[4] * self.factor)
            RZ = round(self.safe_initial_pose[5] * self.factor)
            
            print(f"移动到初始位姿: X={X}, Y={Y}, Z={Z}, RX={RX}, RY={RY}, RZ={RZ}")
            
            # 使用慢速移动到初始位置（速度15，确保安全）
            self.piper.MotionCtrl_2(0x01, 0x00, 15, 0x00)
            self.piper.EndPoseCtrl(X, Y, Z, RX, RY, RZ)
            time.sleep(3)  # 等待移动完成
            
        except Exception as e:
            print(f"移动到初始位姿失败: {e}")
    
    def _wait_motion_done(self):
        """等待机械臂运动完成"""
        start_t = time.time()
        timeout = 5  # 超时时间（秒）
        
        while self._is_in_motion():
            time.sleep(0.01)  # 10ms检查间隔
            if time.time() - start_t > timeout:
                print("等待运动完成超时")
                return False
        return True
    
    def _is_in_motion(self):
        """检查机械臂是否在运动"""
        try:
            return bool(self.piper.GetArmStatus().arm_status.motion_status)
        except:
            return False
    
    def _get_current_position(self):
        """获取当前机械臂末端位姿"""
        try:
            # 使用正确的API调用方法获取位姿信息
            endpose = self.piper.GetArmEndPoseMsgs().end_pose
            
            # 转换为毫米和度
            position = [
                endpose.X_axis / self.factor,  # 微米转毫米
                endpose.Y_axis / self.factor,
                endpose.Z_axis / self.factor,
                endpose.RX_axis / self.factor,  # 微度转度
                endpose.RY_axis / self.factor,
                endpose.RZ_axis / self.factor
            ]
            return position
            
        except Exception as e:
            print(f"获取位姿失败: {e}")
            return None
    
    def _move_robot_to_position(self, new_pos, wait_completion=False):
        """移动机械臂到指定位置"""
        try:
            # 只有在需要等待完成时才等待（如重置位置时）
            if wait_completion:
                self._wait_motion_done()
            
            # 转换为微米单位
            X = round(new_pos[0] * self.factor)
            Y = round(new_pos[1] * self.factor)
            Z = round(new_pos[2] * self.factor)
            RX = round(new_pos[3] * self.factor)
            RY = round(new_pos[4] * self.factor)
            RZ = round(new_pos[5] * self.factor)
            
            # 设置运动速度（30Hz控制频率下的连续移动）
            speed = 30  # 使用较快的速度进行连续控制
            self.piper.MotionCtrl_2(0x01, 0x00, speed, 0x00)
            
            # 发送位置控制命令
            self.piper.EndPoseCtrl(X, Y, Z, RX, RY, RZ)
            
            # 更新当前位置
            self.current_position = new_pos.copy()
            
        except Exception as e:
            print(f"移动失败: {e}")
    
    # ==================== 移动控制方法 ====================
    
    def move_x(self, direction):
        """X轴移动（正方向向右，负方向向左）"""
        if self.current_position is None:
            return
        
        new_pos = self.current_position.copy()
        new_pos[0] += direction * self.move_distance_per_frame * 10  # 转换为毫米
        self._move_robot_to_position(new_pos, wait_completion=False)
    
    def move_y(self, direction):
        """Y轴移动（正方向向前，负方向向后）"""
        if self.current_position is None:
            return
        
        new_pos = self.current_position.copy()
        new_pos[1] += direction * self.move_distance_per_frame * 10  # 转换为毫米
        self._move_robot_to_position(new_pos, wait_completion=False)
    
    def move_z(self, direction):
        """Z轴移动（正方向向上，负方向向下）"""
        if self.current_position is None:
            return
        
        new_pos = self.current_position.copy()
        new_pos[2] += direction * self.move_distance_per_frame * 10  # 转换为毫米
        self._move_robot_to_position(new_pos, wait_completion=False)
    
    def reset_to_initial_position(self):
        """重置到初始位置"""
        print("正在重置到初始位置...")
        if self.initial_position is None:
            print("错误: 初始位置未设置，无法重置")
            return
        
        print(f"目标位置: X={self.initial_position[0]:.1f}, Y={self.initial_position[1]:.1f}, Z={self.initial_position[2]:.1f}")
        
        try:
            self._move_robot_to_position(self.initial_position, wait_completion=True)
            print("已重置到初始位置")
        except Exception as e:
            print(f"重置位置失败: {e}")
            import traceback
            traceback.print_exc()
    
    # ==================== 夹爪控制方法 ====================
    
    def open_gripper(self):
        """打开夹爪（100mm宽度）"""
        print("正在打开夹爪...")
        try:
            # 等待运动完成
            self._wait_motion_done()
            # 发送夹爪控制命令 - 100mm宽度
            self.piper.GripperCtrl(100 * 1000, 1000, 0x01, 0)
            # 等待夹爪操作完成
            time.sleep(self.gripper_delay)
            print("夹爪已打开")
        except Exception as e:
            print(f"打开夹爪失败: {e}")
    
    def close_gripper(self):
        """关闭夹爪（完全闭合）"""
        print("正在关闭夹爪...")
        try:
            # 等待运动完成
            self._wait_motion_done()
            # 发送夹爪控制命令 - 0mm宽度（完全闭合）
            self.piper.GripperCtrl(0, 1000, 0x01, 0)
            # 等待夹爪操作完成
            time.sleep(self.gripper_delay)
            print("夹爪已关闭")
        except Exception as e:
            print(f"关闭夹爪失败: {e}")
    
    # ==================== Xbox手柄控制方法 ====================
    
    def handle_joystick_control(self, axes):
        """
        处理摇杆控制 - 持续移动控制
        
        Args:
            axes: 手柄轴数据列表
        """
        if len(axes) < 6:
            return
        
        # 获取摇杆和触发器值
        left_stick_x = -axes[0]  # 左摇杆X轴（取反修正左右方向）
        left_stick_y = -axes[1]  # 左摇杆Y轴（反向）
        lt_trigger = axes[2]     # 左触发器
        rt_trigger = axes[5]     # 右触发器
        
        # 应用死区 - 摇杆死区
        if abs(left_stick_x) < self.deadzone:
            left_stick_x = 0
        if abs(left_stick_y) < self.deadzone:
            left_stick_y = 0
        
        # 触发器特殊处理 - Xbox手柄触发器未按下时为-1，按下时为1
        # 将-1到1的范围映射到0到1的范围，然后应用死区
        if lt_trigger < 0:
            lt_trigger = 0  # 未按下状态设为0
        elif abs(lt_trigger) < self.deadzone:
            lt_trigger = 0  # 死区内设为0
            
        if rt_trigger < 0:
            rt_trigger = 0  # 未按下状态设为0
        elif abs(rt_trigger) < self.deadzone:
            rt_trigger = 0  # 死区内设为0
        
        # 修正轴映射：摇杆X轴控制机械臂Y轴，摇杆Y轴控制机械臂X轴
        # X轴移动控制（左摇杆Y轴）
        if abs(left_stick_y) > self.deadzone:
            self.move_x(left_stick_y)
        
        # Y轴移动控制（左摇杆X轴）
        if abs(left_stick_x) > self.deadzone:
            self.move_y(left_stick_x)
        
        # Z轴移动控制（LT/RT触发器）
        if lt_trigger > 0:
            # LT触发器 - Z轴向下
            self.move_z(-lt_trigger)
        elif rt_trigger > 0:
            # RT触发器 - Z轴向上
            self.move_z(rt_trigger)
        
        # 调试信息 - 显示摇杆状态和移动指令
        if abs(left_stick_x) > 0.01 or abs(left_stick_y) > 0.01 or lt_trigger > 0.01 or rt_trigger > 0.01:
            print(f"摇杆状态: X={left_stick_x:.3f}, Y={left_stick_y:.3f}, LT={lt_trigger:.3f}, RT={rt_trigger:.3f}")
            if abs(left_stick_y) > self.deadzone:
                print(f"  -> 机械臂X轴移动: {left_stick_y * self.move_distance_per_frame * 10:.2f}mm")
            if abs(left_stick_x) > self.deadzone:
                print(f"  -> 机械臂Y轴移动: {left_stick_x * self.move_distance_per_frame * 10:.2f}mm")
            if lt_trigger > 0:
                print(f"  -> Z轴向下: {-lt_trigger * self.move_distance_per_frame * 10:.2f}mm")
            elif rt_trigger > 0:
                print(f"  -> Z轴向上: {rt_trigger * self.move_distance_per_frame * 10:.2f}mm")
        

    def handle_button_control(self, buttons):
        """
        处理按钮控制 - 夹爪和重置
        
        Args:
            buttons: 手柄按钮状态列表
        """
        if len(buttons) < 8:
            current_buttons = buttons + [False] * (8 - len(buttons))
        else:
            current_buttons = buttons[:8]
        
        # 检查按钮状态变化
        for i, button_pressed in enumerate(current_buttons):
            last_pressed = self.last_button_states.get(i, False)
            
            if button_pressed and not last_pressed:  # 按钮被按下
                print(f"按钮 {i} 被按下")
                if i == 4:  # LB按钮 - 关闭夹爪
                    print("LB按钮 - 关闭夹爪")
                    self.close_gripper()
                elif i == 5:  # RB按钮 - 打开夹爪
                    print("RB按钮 - 打开夹爪")
                    self.open_gripper()
                elif i == 7:  # Menu按钮 - 重置位置
                    print("Menu按钮 - 重置位置")
                    self.reset_to_initial_position()
            
            # 更新按钮状态
            self.last_button_states[i] = button_pressed
    
    def read_xbox_inputs(self):
        """读取Xbox手柄输入"""
        # 处理pygame事件
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                return False
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_ESCAPE:
                    return False
        
        # 读取摇杆和按钮状态
        axes = [self.joystick.get_axis(i) for i in range(self.joystick.get_numaxes())]
        buttons = [self.joystick.get_button(i) for i in range(self.joystick.get_numbuttons())]
        
        # 处理摇杆控制
        self.handle_joystick_control(axes)
        
        # 处理按钮控制
        self.handle_button_control(buttons)
        
        return True
    
    def run(self):
        """运行主控制循环"""
        print("\n=============== Piper机械臂Xbox控制系统 V1.0 ===============")
        print("作者: Yinchuan Li")
        print("日期: 2025-01-27")
        print("\n控制说明:")
        print("  [摇杆控制 - 持续移动]")
        print("  - 左摇杆X轴: 控制机械臂Y轴左右移动")
        print("  - 左摇杆Y轴: 控制机械臂X轴前后移动")
        print("  - LT触发器: Z轴向下移动")
        print("  - RT触发器: Z轴向上移动")
        print("  [按钮控制]")
        print("  - LB键: 关闭夹爪")
        print("  - RB键: 打开夹爪")
        print("  - Menu键: 重置到初始位置")
        print("  - ESC键: 退出程序")
        print("=" * 60)
        
        # 主循环
        clock = pygame.time.Clock()
        
        try:
            while True:
                # 读取手柄输入
                if not self.read_xbox_inputs():
                    break
                
                # 控制帧率（30Hz）
                clock.tick(self.control_frequency)
                
        except KeyboardInterrupt:
            print("\n程序被用户中断")
        except Exception as e:
            print(f"程序运行错误: {e}")
            import traceback
            traceback.print_exc()
        finally:
            pygame.quit()
            print("程序已退出")

def main():
    """主函数"""
    try:
        # 初始化机械臂
        piper = C_PiperInterface("can1")
        xbox_controller = XboxPiperController(piper)
        
        # 运行主控制循环
        xbox_controller.run()
        
    except Exception as e:
        print(f"程序初始化错误: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main()
